﻿namespace SpikingVisualisationRLib

open RProvider
open RProvider.RInterop
open RProvider.``base``
open RProvider.graphics
open SpikingAnalyticsLib    // for OneSecondTickDataCollector

[<AbstractClass; Sealed>]
type SpikeVisualisation () =

    /// Show firing events as a spike raster
    // Requires a list or set (a sequence does not work with FSharpChart)
    static member ShowSpikeRaster(firingEvents:#seq<int * int>) =

        let events = Seq.toArray firingEvents
        let timeProjection = fun (time, nindex) -> time
        let nindexProjection = fun (time, nindex) -> nindex

        let fromMillisecond, toMillisecond =
            fst (Seq.minBy timeProjection events), fst (Seq.maxBy timeProjection events)

        let maxNindex = snd (Seq.maxBy nindexProjection events)

        let xData = events |> Seq.map timeProjection
        let yData = events |> Seq.map nindexProjection

        // draw a plot area
        namedParams [
            "x", box [fromMillisecond; toMillisecond];
            "y", box [0; maxNindex];
            "type", box "n";
            "xlab", box "Time (msec)";
            "ylab", box "Neuron Index";
        ]
        |> R.plot |> ignore

        namedParams [
            "x", box xData;
            "y", box yData;
            "type", box "p";
            "pch", box 16;
        ]
        |> R.points |> ignore

    /// Show collected firing data as a spike raster
    static member ShowCollectedData(collector:OneSecondTickDataCollector) =
        let firingEvents =
            collector.AllEventPairs
            |> Seq.toList
        SpikeVisualisation.ShowSpikeRaster(firingEvents)

    /// Show groups of firing events as multiple rows of spike rasters
    static member ShowSpikeRasterGroups(groups:#seq<#seq<int * int>>) =

        let groupList = Seq.toList groups

        // set the margins and number of rows
        namedParams [
            "mar", box [ 2; 4; 1; 4; ];
            "mfrow", box [groupList.Length; 1];
        ]
        |> R.par |> ignore

        groupList
        |> Seq.iter (fun firingEvents -> SpikeVisualisation.ShowSpikeRaster(firingEvents))

    /// Display a pair of spike trains, one above the other
    // Spike trains are represented as boolean arrays of spikes (true -> spike; false -> no spike)
    static member CompareSpikeTrains(spikes0, spikes1, fromMillisecond, toMillisecond) =

        let DrawLine (xDim: float list) (yDim:float list) =
            namedParams [
                "x", box xDim;
                "y", box yDim;
                "type", box "l";
            ]
            |> R.lines |> ignore

        // draw a plot area
        namedParams [
            "x", box [fromMillisecond; toMillisecond];
            "y", box [0.0; 1.75];
            "type", box "n";
            "xaxt",  box "n";
            "yaxt",  box "n";
            "xlab", box "Time (msec)";
            "ylab", box "";
        ]
        |> R.plot |> ignore

        // x-ticks and labels every 100 msecs
        let xpos = [ for x in fromMillisecond..100..toMillisecond -> x ]
        let xposLabels = List.map (fun x -> sprintf "%d" x) xpos |> List.toArray

        // set the x-axis
        namedParams [
            "side", box 1;      // x-axis
            "at", box xpos;
            "labels", box xposLabels;
            "cex.axis", box 1.5;
        ]
        |> R.axis |> ignore

        // draw horizontal baselines
        DrawLine [float fromMillisecond; float toMillisecond] [1.0; 1.0]
        DrawLine [float fromMillisecond; float toMillisecond] [0.25; 0.25]

        // Extract the indices of the spikes
        let ExtractSpikeIndices spikes =
            spikes
            |> Seq.mapi (fun index spike -> if spike then Some(index + fromMillisecond) else None)
            |> Seq.choose id
            |> Seq.toArray

        // Get the indices of all spikes
        let ind0 = ExtractSpikeIndices spikes0
        let ind1 = ExtractSpikeIndices spikes1

        // Draw vertical 'spikes' at each spike index
        for offset in ind0 do
            DrawLine [float offset; float offset] [1.0; 1.5]
        for offset in ind1 do
            DrawLine [float offset; float offset] [0.25; 0.75]

    /// Display a pair of spike trains, one above the other
    // Spike trains are represented as boolean arrays of spikes (true -> spike; false -> no spike)
    static member CompareSpikeTrains(index1, index2, collector:OneSecondTickDataCollector, fromMillisecond, toMillisecond) =

        let EventsToBooleanArray slots spikeList =
            let spikes = Array.zeroCreate<bool> slots
            spikeList
            |> List.iter (fun (time, _) -> spikes.[time - fromMillisecond] <- true)

            spikes

        let spikeList0, spikeList1 =
            collector.AllEventPairs
            |> Seq.filter (fun (time, nindex) -> nindex = index1 || nindex = index2)
            |> Seq.filter (fun (time, nindex) -> time >= fromMillisecond && time <= toMillisecond)
            |> Seq.toList
            |> List.partition (fun (time, nindex) -> nindex = index1)

        let slots = toMillisecond - fromMillisecond + 1
        let spikes0 = EventsToBooleanArray slots spikeList0
        let spikes1 = EventsToBooleanArray slots spikeList1

        SpikeVisualisation.CompareSpikeTrains(spikes0, spikes1, fromMillisecond, toMillisecond)
